"""EEG 数据读取与事件解析（简化版）

功能：
- 读取 BrainVision 三件套（.vhdr/.eeg/.vmrk）
- 提供 EEG 数据矩阵和事件信息
"""

import mne
import numpy as np
from pathlib import Path

class GZMI1Recording:
    """简化的 EEG 数据读取器"""

    def __init__(self, root="D:\code\Python\library_BCI\data\lwx", vhdr_name="lwx_mi_1.vhdr"):
        self.root = Path(root)
        self.vhdr_path = self.root / vhdr_name # 使用传入的 vhdr_name
        
        # 检查文件是否存在
        if not self.vhdr_path.exists():
            raise FileNotFoundError(f"未找到 {self.vhdr_path}")
        
        self.raw = None

    def load(self, preload=True):
        """加载 EEG 数据"""
        self.raw = mne.io.read_raw_brainvision(
            str(self.vhdr_path), preload=preload, verbose=False
        )

    def get_data_matrix(self, channels=None):
        """返回 EEG 数据矩阵 (n_channels, n_times)"""
        if self.raw is None:
            raise RuntimeError("请先调用 load() 方法")
        return self.raw.get_data(picks=channels)

    def get_channel_names(self):
        """返回通道名称列表"""
        if self.raw is None:
            raise RuntimeError("请先调用 load() 方法")
        return self.raw.ch_names

    def get_sampling_rate(self):
        """返回采样率"""
        if self.raw is None:
            raise RuntimeError("请先调用 load() 方法")
        return self.raw.info["sfreq"]

    def get_events(self):
        """提取事件信息"""
        if self.raw is None:
            raise RuntimeError("请先调用 load() 方法")
        
        # 获取事件
        events, event_id = mne.events_from_annotations(self.raw, verbose=False)
        
        # 获取注释信息
        annotations = self.raw.annotations
        
        # 构建事件记录
        records = []
        for i, (onset, duration, desc) in enumerate(zip(annotations.onset, 
                                                        annotations.duration, 
                                                        annotations.description)):
            if desc in event_id:
                onset_sample = int(self.raw.time_as_index(onset)[0])
                records.append({
                    "onset_sample": onset_sample,
                    "onset_time_sec": onset,
                    "description": desc,
                    "code": event_id[desc]
                })
        
        return events, event_id, records

if __name__ == "__main__":
    # 测试代码
    try:
        reader = GZMI1Recording()
        reader.load(preload=True)

        # 打印基本信息
        print(f"采样率: {reader.get_sampling_rate()} Hz")
        print(f"通道数: {len(reader.get_channel_names())}")
        
        # 获取数据矩阵
        data_matrix = reader.get_data_matrix()
        print(f"数据矩阵形状: {data_matrix.shape} (n_channels, n_times)")
        
        # 获取事件
        events, event_id, records = reader.get_events()
        print(f"事件类型: {event_id}")
        print(f"事件总数: {len(records)}")
        
        if records:
            print(f"首个事件: {records[0]}")
            
        print("数据读取成功！")
        
    except Exception as e:
        print(f"错误: {e}")
